import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import cv2
import numpy as np
import torchvision.transforms as transforms

class ImagePro:
    def __init__(self,):
        pass

    # def load_image(self, img_path):
    #     self.img_path = img_path
    #     img = cv2.imread(self.img_path)
    #     img_tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float() / 255.0
    #     return img_tensor

    def load_image(self, img_path, mode=None):
        """
        加载图像指令
        用途：导入本地图像文件
        参数：图像路径
        返回值：结果存入变量
        """
        img = cv2.imread(img_path)

        if mode == 'tensor':
            
            img_tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float() / 255.0
            return img_tensor
        elif mode == 'gray':
            self.img_path = img_path
            img = cv2.resize(img, (80, 80))
            img = img.astype(np.float32) / 255.0
            input_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 转换为灰度图像
            
            return input_image
        else:

            return img_path






    def expand_kernel(self, kernel):
        return torch.cat([kernel, kernel, kernel], dim=1)

    def select_kernel(self, kernel_type):
        if kernel_type == 'sharpen':
            kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]])
        elif kernel_type == 'texture':
            kernel = np.array([[2, -1, -1], [-1, 3, -1], [-1, -1, 2]])
        elif kernel_type == 'dynamic':
            kernel = np.array([[1/3, 1/3, 1/3], [0, 0, 0], [0, 0, 0]])
        elif kernel_type == 'edge':
            kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
        else:
            raise ValueError("Invalid kernel type")
        return kernel

    def apply_convolution(self, img, kernel):
        return F.conv2d(img, kernel, padding=1)

    def apply_pooling(self, feature, pool_type='max'):
        if pool_type == 'max':
            return F.max_pool2d(feature, kernel_size=2, stride=2)
        elif pool_type == 'avg':
            return F.avg_pool2d(feature, kernel_size=2, stride=2)
        else:
            raise ValueError("Unsupported pool type. Use 'max' or 'avg'.")

    def draw(self, img_tensor, feature_map, kernel_type, titleq,title,fontsize,pool_type=''):
 
        plt.rcParams['figure.dpi'] = 300
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        plt.figure(figsize=(3, 1.5))
        plt.subplot(1, 2, 1)
        plt.title(titleq, fontsize=fontsize)
        plt.axis('off')
        img_tensor = img_tensor.squeeze().permute(1, 2, 0).numpy()
        img_rgb = cv2.cvtColor(img_tensor, cv2.COLOR_BGR2RGB)
        plt.imshow(img_rgb, cmap='gray')

        plt.subplot(1, 2, 2)
        plt.axis('off')

        plt.title(title, fontsize=fontsize)
        plt.imshow(feature_map.squeeze())
        plt.show()



    def process_image(self, kernel_type, pool_type='max'):
        kernel = self.select_kernel(kernel_type)
        kernel = torch.from_numpy(kernel).unsqueeze(0).unsqueeze(0).float()
        kernel = self.expand_kernel(kernel)

        feature = self.apply_convolution(self.img_tensor, kernel)
        self.draw(self.img_tensor, feature, kernel_type)

        feature_pooled = self.apply_pooling(feature, pool_type)
        pool_type_name = '-最大池化' if pool_type == 'max' else '-平均池化'
        self.draw(self.img_tensor, feature_pooled, kernel_type, pool_type_name)

    def draw_rgb(self,img_path):
        plt.rcParams['figure.dpi'] = 300
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        image = cv2.imread(img_path)
        b, g, r = cv2.split(image)

        zeros = np.zeros(image.shape[:2], dtype="uint8")
        blue_channel = cv2.merge([b, zeros, zeros])
        green_channel = cv2.merge([zeros, g, zeros])
        red_channel = cv2.merge([zeros, zeros, r])

        plt.figure(figsize=(4, 2))
        plt.subplot(141)
        plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        plt.title('原始图像', fontsize='6')
        plt.axis('off')

        plt.subplot(142)
        plt.imshow(cv2.cvtColor(blue_channel, cv2.COLOR_BGR2RGB))
        plt.title('蓝色通道', fontsize='6')
        plt.axis('off')

        plt.subplot(143)
        plt.imshow(green_channel)
        plt.title('绿色通道', fontsize='6')
        plt.axis('off')

        plt.subplot(144)
        plt.imshow(cv2.cvtColor(red_channel, cv2.COLOR_BGR2RGB))
        plt.title('红色通道', fontsize='6')
        plt.axis('off')

        plt.show()

